#include "common.h"

void test_tensor_operations() {
  pytorch_study::Timer timer("Tensor Operations Test");

  // 1. 创建张量
  std::cout << "1. Creating tensors:" << std::endl;
  auto tensor1 = torch::tensor({1, 2, 3, 4, 5});
  auto tensor2 = torch::rand({2, 3});
  auto tensor3 = torch::zeros({2, 2});

  prtts2(tensor1, "tensor1");
  prtts2(tensor2, "tensor2");
  prtts2(tensor3, "tensor3");

  // 2. 张量运算
  std::cout << "\n2. Tensor operations:" << std::endl;
  auto result1 = tensor1 + tensor1;
  auto result2 = torch::matmul(tensor2, tensor2.transpose(0, 1));

  prtts2(result1, "tensor1 + tensor1");
  prtts2(result2, "matmul result");

  // 3. 张量索引和切片
  std::cout << "\n3. Tensor indexing:" << std::endl;
  auto large_tensor = torch::rand({5, 5});
  std::cout << "Original tensor:\n" << large_tensor << std::endl;
  std::cout << "First row:\n" << large_tensor[0] << std::endl;
  std::cout << "2x2 slice:\n"
            << large_tensor.index(
                   {torch::indexing::Slice(0, 2), torch::indexing::Slice(0, 2)})
            << std::endl;
}